import matplotlib.pyplot as plt
import matplotlib.animation as animation
import numpy as np
import time


# 修改后的计数排序算法
def counting_sort(arr):
    n = len(arr)
    steps = []
    # 找到数组中的最大值和最小值
    max_val = max(arr)
    min_val = min(arr)

    # 初始化计数数组
    count = [0] * (max_val - min_val + 1)

    # 计数每个元素的频率
    for i in range(n):
        count[arr[i] - min_val] += 1
        steps.append((arr.copy(), i, arr[i]))  # 记录排序过程和当前处理的元素

    # 重建排序数组
    index = 0
    for i in range(len(count)):
        while count[i] > 0:
            arr[index] = i + min_val
            count[i] -= 1
            index += 1
            steps.append((arr.copy(), index - 1, i + min_val))  # 记录排序过程

    return steps


# 随机生成一个数组
arr = np.random.randint(0, 100, 20)

# 获取排序步骤
steps = counting_sort(arr.copy())

# 创建动画
fig, ax = plt.subplots()
ax.set_xlim(-1, len(arr))
ax.set_ylim(0, max(arr) + 10)

bar_rects = ax.bar(range(len(arr)), arr, align="edge")


def animate(i):
    current_step = steps[i][0]
    highlight_idx = steps[i][1]
    highlight_val = steps[i][2]

    for idx, (rect, val) in enumerate(zip(bar_rects, current_step)):
        if idx == highlight_idx:
            rect.set_color('r')  # 将处理的柱子设为红色
        else:
            rect.set_color('b')  # 其余柱子设为蓝色
        rect.set_height(val)
    time.sleep(0.05)  # 可以适当调整延迟以观察排序过程

    return bar_rects


ani = animation.FuncAnimation(fig, animate, frames=len(steps), interval=50, blit=True)

plt.show()
